//===----------------------------------------------------------------------===//
//
// Copyright (C) 2022 Sophgo Technologies Inc.  All rights reserved.
//
// TPU-MLIR is licensed under the 2-Clause BSD License except for the
// third-party components.
//
//===----------------------------------------------------------------------===//

#include "tpu_mlir/Backend/CV18xx/CV18xx.h"
#include "tpu_mlir/Support/MathUtils.h"
#include <valarray>

using namespace tpu_mlir::backend;

LogicalResult tpu::ShapeTileOp::init(InferenceParameter &p) {
  return success();
}
void tpu::ShapeTileOp::deinit(InferenceParameter &p) {}

LogicalResult tpu::ShapeTileOp::inference(InferenceParameter &p) {
  auto in_shape = module::getShape(getInput());
  auto tmp_size = in_shape.size();
  std::vector<int64_t> tile_vec;
  if (!getTileT()) {
    auto tile_v = module::getI64Array(getTile());
    tile_vec = *tile_v;
  } else {
    for (int i = 0; i < tmp_size; ++i) {
      tile_vec.emplace_back((int64_t)p.inputs[1][i]);
    }
  }

  std::vector<int64_t> out_shape(in_shape.size());
  std::transform(tile_vec.begin(), tile_vec.end(), in_shape.begin(),
                 out_shape.begin(), [](int a, int b) { return a * b; });
  module::setShape(getOutput(), out_shape);

  int last_i = tile_vec.size() - 1;
  for (int i = 0; i < tile_vec.size(); ++i) {
    last_i = tile_vec.size() - i - 1;
    if (tile_vec[last_i] != 1)
      break;
  }

  auto last_op = p.inputs[0];
  std::vector<int64_t> tmp_shape(in_shape.begin(), in_shape.end());

  for (int i = 0; i < last_i + 1; ++i) {
    if (tile_vec[i] == 1)
      continue;
    int len = std::accumulate(tmp_shape.begin(), tmp_shape.end(), 1,
                              std::multiplies<int64_t>());
    float *cur_input = new float[len];
    std::copy(last_op, last_op + len, cur_input);

    function_tile(cur_input, p.outputs[0],
                  llvm::ArrayRef<int64_t>(tmp_shape.data(), tmp_size), i,
                  tile_vec[i]);
    last_op = p.outputs[0];
    tmp_shape[i] *= tile_vec[i];
    delete[] cur_input;
  }
  return success();
}

bool tpu::ShapeTileOp::support_multi_core() { return false; }
